#include "stdafx.h"
#include "Volume.h"
#include "DriveStream.h"

using namespace System;
using namespace IO;

StringList Volume::getPaths(LPCWSTR volumeName)
{
	StringList result;

    DWORD  count		= MAX_PATH + 1;
    LPWSTR paths		= NULL;
    PWCHAR pathIdx		= NULL;
    bool   success		= false;

    for (;;) 
    {
        paths = (LPWSTR) new BYTE [count * sizeof(WCHAR)];

        if ( !paths ) 
        {
            return result;
        }
        success = !!GetVolumePathNamesForVolumeNameW(
            volumeName, paths, count, &count
            );

        if ( success ) 
        {
            break;
        }

        if ( GetLastError() != ERROR_MORE_DATA ) 
        {
            break;
        }

        //  Try again with the
        //  new suggested size.
        
		delete [] paths;
        paths = NULL;
    }

    if ( success )
    {
        //
        //  Display the various paths.
        for ( pathIdx = paths; 
              pathIdx[0] != L'\0'; 
              pathIdx += wcslen(pathIdx) + 1 ) 
        {
            //wprintf(L"  %s", pathIdx);

			result.push_back(String(pathIdx));
        }
    }

    if ( paths != NULL ) 
    {
        delete [] paths;
        paths = NULL;
    }

    return result;
}

VolumeList Volume::getVolumes()
{
	VolumeList result;

    DWORD  count				= 0;
    WCHAR  deviceName[MAX_PATH] = L"";
    HANDLE FindHandle           = INVALID_HANDLE_VALUE;
    size_t Index                = 0;
    BOOL   success              = FALSE;
    WCHAR  volumeName[MAX_PATH] = L"";

	WCHAR 	userDefinedName[MAX_PATH] = L"";
	DWORD 	nameSize = ARRAYSIZE(userDefinedName);
	WCHAR 	fileSystem[MAX_PATH];
	DWORD 	fileSystemSize = ARRAYSIZE(fileSystem);

    FindHandle = FindFirstVolumeW(volumeName, ARRAYSIZE(volumeName));

    if (FindHandle == INVALID_HANDLE_VALUE)
    {
        DWORD err = GetLastError();
		WCHAR	errMsg[256];
		wsprintfW(errMsg, L"FindFirstVolumeW failed with error code %d\n", err);
		
		throw (Exception::ExceptionBase(errMsg));
    }

    for (;;)
    {
		UDVolumeInformation info;
		memset(deviceName, 0, sizeof(deviceName));
		memset(userDefinedName, 0, sizeof(userDefinedName));
		memset(fileSystem, 0, sizeof(fileSystem));
        //
        //  Skip the \\?\ prefix and remove the trailing backslash.
        Index = wcslen(volumeName) - 1;

        if (volumeName[0]     != L'\\' ||
            volumeName[1]     != L'\\' ||
            volumeName[2]     != L'?'  ||
            volumeName[3]     != L'\\' ||
            volumeName[Index] != L'\\') 
        {
			WCHAR	errMsg[256];
			wsprintfW(errMsg, L"FindFirstVolumeW/FindNextVolumeW returned a bad path: %s\n", volumeName);
			
			throw (Exception::ExceptionBase(errMsg));
        }

        //
        //  QueryDosDeviceW does not allow a trailing backslash,
        //  so temporarily remove it.
        volumeName[Index] = L'\0';
        count = QueryDosDeviceW(&volumeName[4], deviceName, ARRAYSIZE(deviceName)); 
        volumeName[Index] = L'\\';

        if ( count == 0 ) 
        {
			DWORD err = GetLastError();
			WCHAR	errMsg[256];
			wsprintfW(errMsg, L"QueryDosDeviceW failed with error code %d\n", err);
			
			throw (Exception::ExceptionBase(errMsg));
        }
		if (!GetVolumeInformationW(
				volumeName,
				userDefinedName,
  				nameSize,
  				&info.serialNumber,
  				&info.maxNameLen,
  				&info.flags,
  				fileSystem,
  				fileSystemSize
			))
		{
			DWORD	err = GetLastError();
			WCHAR	errMsg[256];
			wsprintfW(errMsg, L"GetVolumeInformationW with error code %d\n", err);
			//only warning here ...
			//throw (Exception::ExceptionBase(errMsg));
		}
        
		StringList pathList = getPaths(volumeName);
		info.deviceName = deviceName;
		info.path = pathList;
		info.volumeName = volumeName;
		info.fileSystem = fileSystem;
		info.userDefinedName = userDefinedName;
		result.push_back(info);
        //
        //  Move on to the next volume.
        if (!FindNextVolumeW(FindHandle, volumeName, ARRAYSIZE(volumeName)))
            break;
    }

    FindVolumeClose(FindHandle);
    FindHandle = INVALID_HANDLE_VALUE;
	
	return result;
}

PartitionList	Volume::getAllPartitions(VolumeParsersList parsers)
{
	PartitionList result;
	
	DriveList	drives  = getPartitions();	
	VolumeList  volumes = getVolumes();
	
	for (DriveI driveI = drives.begin(); driveI != drives.end(); ++ driveI)
	{
		PartitionInfo crtPartitionInfo = {0, };

		DWORD serialNumber = 0;
		switch (driveI->second.type)
		{
		case NTFS_PARTITION :
			serialNumber = parsers[driveI->second.type]->getSerialNumber(driveI->second);	
			break;
		default:
			//throw (Exception::ExceptionBase(L"Unreacheable code exception"));
			break;
		}
		
		for (VolumeI volI = volumes.begin(); volI != volumes.end(); ++volI)
		{
			if (volI->serialNumber == serialNumber)
			{
				crtPartitionInfo.rawInfo = driveI->second;
				crtPartitionInfo.systemInfo = *volI;
				result.push_back (crtPartitionInfo);
				break;
			}
		}
	}
	return result;
}

DriveList		Volume::getPartitions(void)
{
	int i;
	DWORD  bytes;
	DriveList   drives;
	PARTITION	*PartitionTable;
	DriveInfo		stDrive;

	BYTE *mbrSector;
	WORD nDrive =0;

	char buf[64];
	char text[255];

	DWORD mainPrevRelSector =0;
	DWORD prevRelSector =0;

	//HANDLE hDrive = CreateFile(PHYSICAL_DRIVE, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, 0, OPEN_EXISTING, 0, 0);
	//if(hDrive == INVALID_HANDLE_VALUE)
	//	throw Exception::ExceptionBase(L"Error parsing partitions");

	//nRet = ReadFile(hDrive, mbrSector, 512, &bytes, 0);
	//if(!nRet)
	//	throw Exception::ExceptionBase(L"Error parsing partitions");

	DriveInfo dummy = {0, };
	DriveStream		driveStm(dummy, NONE);
	MemoryBuffer	mem(0x400);
	driveStm.read(mem, 0, 512);
	mbrSector = mem.getData();
	prevRelSector =0;
	mainPrevRelSector =0;
/*
	There is room for only four (4) entries in each Partition Table.
	They begin at 01BE, 01CE, 01DE and 01EE and are each 16 bytes long.
*/
	//given that, we will parse all the 16 bytes there and fill 
	//a new drive witch we add to the list of drives we already have
	
	PartitionTable = (PARTITION*)(mbrSector + 0x1BE);

	for(i=0; i<4; i++) 
	{
		stDrive.wCylinder = PartitionTable->cylinder;
		stDrive.head = PartitionTable->head;
		stDrive.sector = PartitionTable->sector;
		stDrive.nSectors = PartitionTable->numberOfSectors;
		stDrive.type = ((PartitionTable->type == PART_EXTENDED) || (PartitionTable->type == PART_DOSX13X)) ? EXTENDED_PART:BOOT_RECORD;

		if((PartitionTable->type == PART_EXTENDED) || (PartitionTable->type == PART_DOSX13X))
		{
			mainPrevRelSector			=		PartitionTable->firstPartitionSector;
			stDrive.NTRelativeSector	=		mainPrevRelSector;
		}
		else
		{
			stDrive.NTRelativeSector = mainPrevRelSector + PartitionTable->firstPartitionSector;
		}

		if(stDrive.type == EXTENDED_PART)
			break;

		if(PartitionTable->type == 0)
			break;
		//I do this for information purpouses only
		//I am only implementing NTFS
#pragma warning( disable: 4996 )
		stDrive.type = PartitionTable->type;
		switch(PartitionTable->type)
		{
		case PART_DOS2_FAT:
			strcpy_s(buf, _countof(buf),"FAT12");
			break;
		case PART_DOSX13:
		case PART_DOS4_FAT:
		case PART_DOS3_FAT:
			strcpy(buf, "FAT16");
			break;
		case PART_DOS32X:
		case PART_DOS32:
			strcpy(buf, "FAT32");	 //Normal FAT32
			break;
		case NTFS_PARTITION:
			strcpy(buf, "NTFS");	// NTFS
			break;
		case UNIX_PARTITION:
			strcpy(buf, "UNIX");	
			break;
		case LINUX_PARTITION:
			strcpy(buf, "LINUX");	
			break;
		case LINUX_SWAP_PARTITION:
			strcpy(buf, "LINUX_SWAP");	
			break;
		default:
			strcpy(buf, "Unknown");
			break;
		}
		wsprintf(text, "%s Drive %d", buf, nDrive);
		printf("%s\n", text);
		drives[text]	=	stDrive;
		PartitionTable++;
		nDrive++;
	}

	if(i==4)
		return drives;
	//same logic for extended partitions
	for(int LogiHard=0; LogiHard<50; LogiHard++) // scanning extended partitions
	{
		if(stDrive.type == EXTENDED_PART)
		{
			LARGE_INTEGER n64Pos;

			n64Pos.QuadPart = ((LONGLONG) stDrive.NTRelativeSector) * 512;
			
			driveStm.seek(n64Pos.QuadPart);
			//nRet = SetFilePointer(hDrive, n64Pos.LowPart,&n64Pos.HighPart, FILE_BEGIN);
			//if(nRet == 0xffffffff)
			//	throw Exception::ExceptionBase(L"Error parsing partitions");

			bytes = 0;

			//nRet = ReadFile(hDrive, mbrSector, 512, (DWORD *) &bytes, NULL);
			driveStm.read(mem, 0, 512);
			mbrSector = mem.getData();
			
			PartitionTable = (PARTITION *) (mbrSector+0x1BE);

			for(i=0; i<4; i++)
			{
				stDrive.wCylinder = PartitionTable->cylinder;
				stDrive.head = PartitionTable->head;
				stDrive.nSectors = PartitionTable->numberOfSectors;
				stDrive.sector = PartitionTable->sector;
				stDrive.firstPartitionSector = 0;
				stDrive.type = ((PartitionTable->type == PART_EXTENDED) || (PartitionTable->type == PART_DOSX13X)) ? EXTENDED_PART:BOOT_RECORD;
				
				if((PartitionTable->type == PART_EXTENDED) || (PartitionTable->type == PART_DOSX13X))
				{
					prevRelSector = PartitionTable->firstPartitionSector;
					stDrive.NTRelativeSector = prevRelSector + mainPrevRelSector;
				}
				else
				{
					stDrive.NTRelativeSector = mainPrevRelSector + prevRelSector + PartitionTable->firstPartitionSector;
				}

				if(stDrive.type == EXTENDED_PART)
					break;

				if(PartitionTable->type == 0)
					break;

				switch(PartitionTable->type)
				{
				case PART_DOS2_FAT:
					strcpy(buf, "FAT12");
					break;
				case PART_DOSX13:
				case PART_DOS4_FAT:
				case PART_DOS3_FAT:
					strcpy(buf, "FAT16");
					break;
				case PART_DOS32X:
				case PART_DOS32:
					strcpy(buf, "FAT32");			//Normal FAT32
					break;
				case NTFS_PARTITION:
					stDrive.type = NTFS_PARTITION;
					strcpy(buf, "NTFS");
					break;
				case UNIX_PARTITION:
					strcpy(buf, "UNIX");	
					break;
				case LINUX_PARTITION:
					strcpy(buf, "LINUX");	
					break;
				case LINUX_SWAP_PARTITION:
					strcpy(buf, "LINUX_SWAP");	
					break;
				default:
					strcpy(buf, "Unknown");
					break;
				}
				wsprintf(text, "%s Drive %d", buf, nDrive);
				drives[text]	=	stDrive;
				printf("%s\n", text);
				PartitionTable++;
				nDrive++;
			}
			if(i==4)
				break;
		}
	}

	return drives;
}
